# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/18 15:15 
@Author  : ZhangShenao 
@File    : 3.使用SQLAgent执行数据库操作.py
@Desc    : 使用SQLAgent执行数据库操作
"""
import os

import dotenv
from langchain.agents import AgentType
from langchain_community.agent_toolkits import create_sql_agent, SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI

# 加载环境变量
dotenv.load_dotenv()

# 连接到本地SQLite数据库
db = SQLDatabase.from_uri(database_uri='sqlite:///flower_shop.db')

# 创建LLM
llm = ChatOpenAI(model_name='gpt-3.5-turbo',
                 openai_api_base=os.getenv('OPENAI_API_BASE'),
                 temperature=0.,  # 这里设置temperature=0,意味着模型的输出会更加确定
                 verbose=True,
                 )

# 创建SQLAgent,它可以驱动LLM来执行数据库相关操作
sql_agent = create_sql_agent(
    llm=llm,
    toolkit=SQLDatabaseToolkit(db=db, llm=llm),  # 使用数据库工具箱SQLDatabaseToolkit
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
)

# 使用Agent来执行数据库查询操作
questions = ['哪种鲜花的存货数量最少？', '鲜花的平均销售价格是多少？']
for question in questions:
    print(sql_agent.run(question))
